package com.edata.bigdata.utils

import org.apache.spark.graphx.{Edge, Graph}
import org.apache.spark.sql.{DataFrame, Row}

object GraphUtils {

  def createGraphXByDataFrame(vertex: DataFrame,
                              edge: DataFrame,
                              vIdField: String,
                              srcIdField: String,
                              dstIdField: String
                             ): Graph[Row, Row] = {

    val vertexRDD = vertex.rdd.map(row => {
      val vtxFieldInd = row.fieldIndex(vIdField)
      val vtxFieldType = row.schema.fields(vtxFieldInd).dataType.typeName
      val vertexId = DataFrameUtils.getValueByCol(row,vIdField,vtxFieldType).toLong
      (vertexId, row)
    })
    val edgeRDD = edge.rdd.map(row => {
      val srcFieldInd = row.fieldIndex(srcIdField)
      val srcFieldType = row.schema.fields(srcFieldInd).dataType.typeName
      var srcId = DataFrameUtils.getValueByCol(row,srcIdField,srcFieldType).toLong
      val dstFieldInd = row.fieldIndex(dstIdField)
      val dstFieldType = row.schema.fields(dstFieldInd).dataType.typeName
      val dstId = DataFrameUtils.getValueByCol(row,dstIdField,dstFieldType).toLong
      Edge(srcId, dstId, row)
    })
    Graph(vertexRDD, edgeRDD)
  }
}
